#!/usr/bin/env python3
import asyncio
import json
import os
import signal
import sys,argparse
import websockets
import random
#websocket随机用户昵称
name = {}
#频道广播
users = {}
#rooms 查询用户房间和websocket
rooms = {}
#历史记录
history = {}
parser = argparse.ArgumentParser()
parser.add_argument('-p', '--port', type=int, help='port number used for start',required=True)
args = parser.parse_args()

def signal_handler(signal, frame):
    sys.exit(0)


async def error(websocket, message):
    """
    Send an error message.

    """
    event = {
        "type": "error",
        "message": message,
    }
    await websocket.send(json.dumps(event))


async def echo(websocket):
    """
    收发消息
    """
    async for message in websocket:
        event = json.loads(message)
        #根据api-token查找房间号
        room_id = name[websocket]['room_id']
        #历史记录数量
        if "show_num" in event["type"]:
            r = {
                "type": "show_num",
                "content": "{}"
            }
            if room_id in history : 
                h = len(history[room_id])
            else :
                h = 0
            r['content'] = r['content'].format(h)
            message = json.dumps(r,ensure_ascii=False)
            await websocket.send(message)
        #历史记录
        if "history" in event["type"]:
            start = int(event['page'])
            end = int(event['page_num'])
            num = int(event['num'])
            if room_id in history : 
                n = history[room_id][:num]
                h = n[start*end - end:start*end]
            else :
                h = []
            r = {'type':'history','content':h}
            message = json.dumps(r,ensure_ascii=False)
            await websocket.send(message)
        #聊天
        if "chat" in event['type']:
            chat = {
                "type": "chat",
                "username": "{}",
                "content": "{}",
                "icon":"{}"
            }
            if not "username" in event :
                chat['username'] = name[websocket]['username']
            if "username" in event:
                chat["username"] = event["username"].format(event["username"])
            if "content" in event :
                chat["content"] = event["content"].format(event["content"])
            if "icon" in event :
                chat["icon"] = event["icon"].format(event["icon"])
            #加入历史记录
            if not room_id in history:
                 history[room_id] = []
            history[room_id].append(chat)
            message = json.dumps(chat,ensure_ascii=False)
            websockets.broadcast(users[room_id], message)
        #点赞
        if "good" in event['type']:
            good = {
                "type": "good",
                "num":"{}"
            }
            if "num" in event:
                #print(rooms[event["api_token"]])
                if not room_id in rooms:
                    rooms[room_id] = 1
                else :
                    rooms[room_id] = rooms[room_id] + 1
                good["num"] = rooms[room_id]
            message = json.dumps(good,ensure_ascii=False)
            websockets.broadcast(users[room_id], message)
        #心跳支持
        if "heartbeat" in event["type"]:
            heartbeat = {"type":"heartbeat","content":"pong"}
            message = json.dumps(heartbeat,ensure_ascii=False)
            await websocket.send(message)
        # await websocket.wait_closed()


async def start(websocket,data):
    """
    认证检查api_token是否合法不合法关闭链接
    {"type":"connect","game_id":"20201","username":"强哥","game_name":"某某比赛","api_token":"app登录后返回的token"}
    """
    #认证
    if not "api_token" in data:
        websocket.close()
    api = data["api_token"]
    #用户昵称不存在随机一个
    if not 'username' in data :
        data['username'] = 'running派_{}'.format(random.randint(10000,99999))
    name[websocket] = {'username':data['username'],'room_id':data["game_id"]}
    rooms[api] = {"room":data["game_id"],"websocket":websocket,'username':data['username']}
    if not data["game_id"] in users :
         users[data["game_id"]] = []
    users[data["game_id"]].append(websocket)
    #print(users[data["game_id"]],rooms)
    try:
        #公告
        notice = {
        "type": "notice",
        "username": "官方",
        "content": "欢迎进入{}赛事直播间，实时查看选手位置、榜单、进程、成绩，并为自己喜欢的大神助力打call! 请理性打赏，严禁未成年人直播或打赏。"
        }
        notice["content"] = notice["content"].format(data["game_name"])
        message = json.dumps(notice,ensure_ascii=False)
        await websocket.send(message)
        #加入房间全局广播
        join = {
            "type": "join",
            "username": "{}",
            "content": "进入直播间"
        }
        join["username"] = join["username"].format(data["username"])
        message = json.dumps(join,ensure_ascii=False)
        websockets.broadcast(users[data["game_id"]], message)
        #点赞
        good = {
            "type": "good",
            "num": "{}"
        }
        if data["game_id"] in rooms:
            num = rooms[data["game_id"]]
        else :
            num = 10000
        good["num"] = good["num"].format(num)
        message = json.dumps(good,ensure_ascii=False)
        websockets.broadcast(users[data["game_id"]], message)
        #榜单头像
        if "rank" in rooms:
            if data["game_id"] in rooms["rank"]:
                head_img = rooms["rank"][data["game_id"]]
                message = json.dumps(head_img,ensure_ascii=False)
                websockets.broadcast(users[data["game_id"]], message)
        #聊天，点赞，心跳，查看历史记录
        await echo(websocket)
    except websockets.exceptions.ConnectionClosedError as e:
        pass
    finally:
        del rooms[api]
        users[data["game_id"]].remove(websocket)

#排行
async def rank(websocket,data):
    """
    排行榜

    """
    try:
        message = json.dumps(data,ensure_ascii=False)
        websockets.broadcast(users[data["game_id"]], message)
    finally:
       pass

#礼物
async def gift(websocket,data):
    """
    礼物

    """
    try:
        message = json.dumps(data,ensure_ascii=False)
        websockets.broadcast(users[data["game_id"]], message)
    finally:
       pass
#榜单头像
async def head_img(websocket,data):
    """
    榜单头像
    rooms[20201]['rank']
    """
    try:
        if not "rank" in rooms:
            rooms["rank"] = {}
            rooms["rank"][data["game_id"]] = data
        else :
            rooms["rank"][data["game_id"]] = data
        message = json.dumps(data,ensure_ascii=False)
        websockets.broadcast(users[data["game_id"]], message)
    finally:
       pass

async def handler(websocket):
    """
    处理认证和推送

    {"type":"connect","game_id":"20201","username":"强哥","game_name":"某某比赛","api_token":"app登录后返回的token"}

    """
    message = await websocket.recv()
    event = json.loads(message)
    if "rank" in event['type']:
        await rank(websocket, event)
    elif "gift" in event['type']:
        await gift(websocket,event)
    elif "head_img" in event['type']:
        await head_img(websocket,event)
    elif "connect" in event['type']:
        await start(websocket,event)
    else :
        pass


async def main():
    # 主服务
    loop = asyncio.get_running_loop()
    stop = loop.create_future()
    loop.add_signal_handler(signal.SIGTERM, stop.set_result, None)
    port = int(args.port)
    async with websockets.serve(handler, "0.0.0.0", port):
        await stop
        


if __name__ == "__main__":
    signal.signal(signal.SIGINT, signal_handler)
    asyncio.run(main())